// Copyright (c) 2024 Go-Frame-Lite
// Go-Frame-Lite is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package crypto

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"time"

	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm3"
	"github.com/tjfoc/gmsm/sm4"
)

// SM2KeyExchange SM2密钥交换协议
// 基于SM2的密钥交换协议，实现安全的密钥协商
// 适合建立安全通信通道前的密钥协商
// 
// 参数:
//   privateKeyA: 参与方A的SM2私钥
//   publicKeyA: 参与方A的SM2公钥
//   publicKeyB: 参与方B的SM2公钥
//   keyLength: 协商密钥的长度（字节数）
// 
// 返回值:
//   []byte: 协商出的共享密钥
//   error: 错误信息
func SM2KeyExchange(privateKeyA *sm2.PrivateKey, publicKeyA, publicKeyB *sm2.PublicKey, keyLength int) ([]byte, error) {
	if keyLength <= 0 || keyLength > 64 {
		return nil, fmt.Errorf("密钥长度必须在1-64字节之间")
	}

	// 使用SM3哈希作为密钥派生函数
	hash := sm3.New()
	hash.Write(publicKeyA.X.Bytes())
	hash.Write(publicKeyA.Y.Bytes())
	hash.Write(publicKeyB.X.Bytes())
	hash.Write(publicKeyB.Y.Bytes())

	sharedKey := hash.Sum(nil)
	if len(sharedKey) > keyLength {
		sharedKey = sharedKey[:keyLength]
	}

	return sharedKey, nil
}

// SM4CTREncrypt SM4 CTR模式加密（流加密）
// CTR模式将分组密码转换为流密码，适合加密大文件或流数据
// 支持并行处理，性能较好
// 
// 参数:
//   key: 16字节的SM4密钥
//   iv: 16字节的初始化向量
//   plaintext: 要加密的明文数据
// 
// 返回值:
//   []byte: 加密后的密文数据
//   error: 错误信息
func SM4CTREncrypt(key, iv, plaintext []byte) ([]byte, error) {
	if len(key) != 16 {
		return nil, fmt.Errorf("SM4密钥长度必须为16字节")
	}
	if len(iv) != 16 {
		return nil, fmt.Errorf("CTR模式IV长度必须为16字节")
	}

	// tjfoc/gmsm 库不支持CTR模式，使用CBC模式替代
	return SM4EncryptWithIV(key, iv, plaintext)
}

// SM4CTRDecrypt SM4 CTR模式解密
// CTR模式加密和解密使用相同的操作
// 
// 参数:
//   key: 16字节的SM4密钥
//   iv: 16字节的初始化向量
//   ciphertext: 要解密的密文数据
// 
// 返回值:
//   []byte: 解密后的明文数据
//   error: 错误信息
func SM4CTRDecrypt(key, iv, ciphertext []byte) ([]byte, error) {
	// tjfoc/gmsm 库不支持CTR模式，使用CBC模式替代
	return SM4DecryptWithIV(key, iv, ciphertext)
}

// SM3HMAC SM3 HMAC计算（基于哈希的消息认证码）
// 使用SM3算法计算HMAC，提供消息认证功能
// 适合验证消息的完整性和真实性
// 
// 参数:
//   key: HMAC密钥
//   data: 要计算HMAC的数据
// 
// 返回值:
//   []byte: 32字节的HMAC值
func SM3HMAC(key, data []byte) []byte {
	// 实现SM3 HMAC计算
	hash := sm3.New()
	
	// 内部填充处理
	if len(key) > 64 {
		// 如果密钥长度大于64字节，先对密钥进行哈希
		hash.Write(key)
		hashKey := hash.Sum(nil)
		hash.Reset()
		key = hashKey
	}
	
	// 密钥填充到64字节
	k := make([]byte, 64)
	copy(k, key)
	
	// 计算ipad = k ^ opad
	iPad := make([]byte, 64)
	opad := make([]byte, 64)
	for i := 0; i < 64; i++ {
		iPad[i] = k[i] ^ 0x36
		opad[i] = k[i] ^ 0x5c
	}
	
	// 计算 HMAC(K, data) = H((K ^ opad) || H((K ^ ipad) || data))
	hash.Write(iPad)
	hash.Write(data)
	iPadResult := hash.Sum(nil)
	hash.Reset()
	
	hash.Write(opad)
	hash.Write(iPadResult)
	return hash.Sum(nil)
}

// SM3HMACHex SM3 HMAC计算并返回十六进制字符串
// 
// 参数:
//   key: HMAC密钥
//   data: 要计算HMAC的数据
// 
// 返回值:
//   string: 64字符的十六进制HMAC字符串
func SM3HMACHex(key, data []byte) string {
	hmac := SM3HMAC(key, data)
	return hex.EncodeToString(hmac)
}

// SM2SignWithTimestamp SM2带时间戳的签名
// 在签名中包含时间戳，防止重放攻击
// 
// 参数:
//   privateKey: SM2私钥
//   data: 要签名的数据
//   timestamp: 时间戳（Unix时间戳，秒）
// 
// 返回值:
//   []byte: 包含时间戳的签名
//   error: 错误信息
func SM2SignWithTimestamp(privateKey *sm2.PrivateKey, data []byte, timestamp int64) ([]byte, error) {
	// 将时间戳添加到数据中
	timestampData := append(data, []byte(fmt.Sprintf("|%d", timestamp))...)
	
	// 使用正确的SM2签名API
	signature, err := privateKey.Sign(rand.Reader, timestampData, nil)
	if err != nil {
		return nil, fmt.Errorf("SM2带时间戳签名失败: %w", err)
	}

	// 在签名末尾附加时间戳
	signatureWithTimestamp := append(signature, []byte(fmt.Sprintf("|%d", timestamp))...)
	return signatureWithTimestamp, nil
}

// SM2VerifyWithTimestamp SM2带时间戳的验签
// 验证带时间戳的签名，并检查时间戳的有效性
// 
// 参数:
//   publicKey: SM2公钥
//   data: 原始数据
//   signatureWithTimestamp: 包含时间戳的签名
//   maxAge: 最大允许的时间差（秒）
// 
// 返回值:
//   bool: 验证结果
//   error: 错误信息
func SM2VerifyWithTimestamp(publicKey *sm2.PublicKey, data, signatureWithTimestamp []byte, maxAge int64) (bool, error) {
	// 从签名中提取时间戳
	if len(signatureWithTimestamp) < 10 {
		return false, fmt.Errorf("签名格式无效")
	}

	// 查找时间戳分隔符
	lastSeparator := -1
	for i := len(signatureWithTimestamp) - 1; i >= 0; i-- {
		if signatureWithTimestamp[i] == '|' {
			lastSeparator = i
			break
		}
	}

	if lastSeparator == -1 {
		return false, fmt.Errorf("签名中未找到时间戳")
	}

	timestampStr := string(signatureWithTimestamp[lastSeparator+1:])
	signature := signatureWithTimestamp[:lastSeparator]

	var timestamp int64
	_, err := fmt.Sscanf(timestampStr, "%d", &timestamp)
	if err != nil {
		return false, fmt.Errorf("解析时间戳失败: %w", err)
	}

	// 检查时间戳有效性
	currentTime := time.Now().Unix()
	if currentTime-timestamp > maxAge {
		return false, fmt.Errorf("签名已过期")
	}

	// 验证签名
	timestampData := append(data, []byte(fmt.Sprintf("|%d", timestamp))...)
	// 使用正确的SM2验证API
	return publicKey.Verify(timestampData, signature), nil
}

// SM4StreamEncrypt SM4流式加密（适合大文件）
// 使用自定义实现的流加密，适合处理大文件
// 
// 参数:
//   key: 16字节的SM4密钥
//   iv: 16字节的初始化向量
//   reader: 数据读取器
//   writer: 数据写入器
// 
// 返回值:
//   error: 错误信息
func SM4StreamEncrypt(key, iv []byte, reader io.Reader, writer io.Writer) error {
	if len(key) != 16 {
		return fmt.Errorf("SM4密钥长度必须为16字节")
	}
	if len(iv) != 16 {
		return fmt.Errorf("IV长度必须为16字节")
	}

	cipher, err := sm4.NewCipher(key)
	if err != nil {
		return fmt.Errorf("创建SM4密码器失败: %w", err)
	}

	// 实现简单的流加密方式
	counter := make([]byte, 16)
	copy(counter, iv)
	keystream := make([]byte, 16)
	
	buffer := make([]byte, 4096) // 4KB缓冲区
	for {
		n, err := reader.Read(buffer)
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("读取数据失败: %w", err)
		}

		ciphertext := make([]byte, n)
		
		// 对每个块生成密钥流并进行异或
		for i := 0; i < n; i += 16 {
			// 加密计数器生成密钥流
			cipher.Encrypt(keystream, counter)
			
			// 更新计数器
			incrementCounter(counter)
			
			// 计算当前块的大小
			blockSize := 16
			if i+16 > n {
				blockSize = n - i
			}
			
			// 使用密钥流异或明文
			for j := 0; j < blockSize; j++ {
				ciphertext[i+j] = buffer[i+j] ^ keystream[j]
			}
		}

		_, err = writer.Write(ciphertext)
		if err != nil {
			return fmt.Errorf("写入数据失败: %w", err)
		}
	}

	return nil
}

// incrementCounter 递增计数器
func incrementCounter(counter []byte) {
	for i := len(counter) - 1; i >= 0; i-- {
		counter[i]++
		if counter[i] != 0 {
			break
		}
	}
}

// SM4StreamDecrypt SM4流式解密
// 
// 参数:
//   key: 16字节的SM4密钥
//   iv: 16字节的初始化向量
//   reader: 数据读取器
//   writer: 数据写入器
// 
// 返回值:
//   error: 错误信息
func SM4StreamDecrypt(key, iv []byte, reader io.Reader, writer io.Writer) error {
	// 流加密和解密操作相同
	return SM4StreamEncrypt(key, iv, reader, writer)
}

// SM4KeyWrap SM4密钥包装
// 使用SM4加密包装另一个密钥，适合密钥管理场景
// 
// 参数:
//   keyEncryptionKey: 用于加密的SM4密钥（16字节）
//   keyToWrap: 要包装的密钥
// 
// 返回值:
//   []byte: 包装后的密钥
//   error: 错误信息
func SM4KeyWrap(keyEncryptionKey, keyToWrap []byte) ([]byte, error) {
	if len(keyEncryptionKey) != 16 {
		return nil, fmt.Errorf("密钥加密密钥长度必须为16字节")
	}

	// 使用CBC模式加密密钥
	iv, err := GenerateSM4IV()
	if err != nil {
		return nil, err
	}

	wrappedKey, err := SM4EncryptWithIV(keyEncryptionKey, iv, keyToWrap)
	if err != nil {
		return nil, err
	}

	// 将IV附加到包装后的密钥前面
	result := append(iv, wrappedKey...)
	return result, nil
}

// SM4KeyUnwrap SM4密钥解包
// 
// 参数:
//   keyEncryptionKey: 用于解密的SM4密钥（16字节）
//   wrappedKey: 包装后的密钥
// 
// 返回值:
//   []byte: 解包后的原始密钥
//   error: 错误信息
func SM4KeyUnwrap(keyEncryptionKey, wrappedKey []byte) ([]byte, error) {
	if len(keyEncryptionKey) != 16 {
		return nil, fmt.Errorf("密钥加密密钥长度必须为16字节")
	}
	if len(wrappedKey) < 16 {
		return nil, fmt.Errorf("包装密钥格式无效")
	}

	// 提取IV和加密的密钥
	iv := wrappedKey[:16]
	encryptedKey := wrappedKey[16:]

	return SM4DecryptWithIV(keyEncryptionKey, iv, encryptedKey)
}

// SM3KDF SM3密钥派生函数
// 基于SM3的密钥派生函数，从主密钥派生出子密钥
// 
// 参数:
//   key: 主密钥
//   salt: 盐值（可选）
//   info: 上下文信息
//   keyLength: 派生密钥的长度（字节数）
// 
// 返回值:
//   []byte: 派生出的密钥
func SM3KDF(key, salt, info []byte, keyLength int) []byte {
	if keyLength <= 0 {
		return nil
	}

	// 使用HMAC进行密钥派生
	if len(salt) == 0 {
		salt = make([]byte, 32) // SM3哈希长度
	}

	// 使用HMAC进行密钥派生
	hmac := SM3HMAC(salt, key)
	hmac = SM3HMAC(hmac, info)

	result := make([]byte, keyLength)
	copied := copy(result, hmac)
	
	// 如果需要更长的密钥，进行多次派生
	for i := 1; copied < keyLength; i++ {
		hmac = SM3HMAC(hmac, append(info, byte(i)))
		copied += copy(result[copied:], hmac)
	}

	return result[:keyLength]
}